//
// Created by daybeha on 24-1-24.
//
// github :https://github.com/FeiYull/TensorRT-Alpha

#include <filesystem>
#include "utils/yolo.h"
#include "yolov8.h"


using namespace std;

void setParameters(yolo_utils::InitParameter &initParameters) {
  initParameters.class_names = yolo_utils::dataSets::coco80;
  //initParameters.class_names = utils::dataSets::voc20;
  initParameters.num_class = 80; // for coco
  //initParameters.num_class = 20; // for voc2012
  initParameters.batch_size = 1;  // 8;
  //    initParameters.src_w = 1241;
//    initParameters.src_h = 376;
  initParameters.src_w = 1280;
  initParameters.src_h = 560;
  initParameters.dst_h = 640;
  initParameters.dst_w = 640;
  initParameters.input_output_names = {"images", "output0"};
  initParameters.conf_thresh = 0.25f;
  initParameters.iou_thresh = 0.45f;
  initParameters.save_path = "";


  initParameters.is_show = true;
  initParameters.is_save = false;
}

void task(YOLOV8 &yolo, const yolo_utils::InitParameter &param, std::vector<cv::Mat> &imgsBatch, const int &delayTime) {
  yolo_utils::DeviceTimer d_t0;
  yolo.copy(imgsBatch);
  float t0 = d_t0.getUsedTime();
  yolo_utils::DeviceTimer d_t1;
  yolo.preprocess(imgsBatch);
  float t1 = d_t1.getUsedTime();
  yolo_utils::DeviceTimer d_t2;
  yolo.infer();
  float t2 = d_t2.getUsedTime();
  yolo_utils::DeviceTimer d_t3;
  yolo.postprocess(imgsBatch);
  float t3 = d_t3.getUsedTime();

  sample::gLogInfo <<
                   //"copy time = " << t0 / param.batch_size << "; "
                   "preprocess time = " << t1 / param.batch_size << "; "
                                                                    "infer time = " << t2 / param.batch_size << "; "
                                                                                                                "postprocess time = "
                   << t3 / param.batch_size << std::endl;

  yolo_utils::show(yolo.getObjectss(), param.class_names, delayTime, imgsBatch);
//    yolo_utils::save(yolo.getObjectss(), param.class_names, param.save_path, imgsBatch, param.batch_size, batchi);
  yolo.reset();
}

void LoadKittiImages(const string &strPathToSequence,
                     vector<string> &vstrImageLeft, vector<string> &vstrImageRight,
                     vector<double> &vTimestamps)
{
  ifstream fTimes;
  string strPathTimeFile = strPathToSequence + "/times.txt";
  fTimes.open(strPathTimeFile.c_str());
  while(!fTimes.eof())
  {
    string s;
    getline(fTimes,s);
    if(!s.empty())
    {
      stringstream ss;
      ss << s;
      double t;
      ss >> t;
      vTimestamps.push_back(t);
    }
  }

  string strPrefixLeft = strPathToSequence + "/image_0/";
  string strPrefixRight = strPathToSequence + "/image_1/";

  const int nTimes = vTimestamps.size();
  vstrImageLeft.resize(nTimes);
  vstrImageRight.resize(nTimes);

  for(int i=0; i<nTimes; i++)
  {
    stringstream ss;
    ss << setfill('0') << setw(6) << i;
    vstrImageLeft[i] = strPrefixLeft + ss.str() + ".png";
    vstrImageRight[i] = strPrefixRight + ss.str() + ".png";
  }
}

void LoadUrbanIamges(const string &strPathToSequence, vector<string> &vstrImageLeft)
{
  string strPrefixLeft = strPathToSequence + "/";
  for (const auto& entry : std::filesystem::directory_iterator(strPrefixLeft)) {
    if (entry.is_regular_file() && entry.path().extension() == ".png") {
      vstrImageLeft.push_back(entry.path().string());
//        cv::Mat image = cv::imread(entry.path().string());
//        if (image.empty()) {
//          std::cout << "Error: Unable to read image " << entry.path().string() << std::endl;
//        } else {
//          std::cout << "Reading image " << entry.path().string() << std::endl;
////          cv::imshow("Image", image);
////          cv::waitKey(0);
//        }
    }
  }
}

int main(int argc, char** argv)
{
  yolo_utils::InitParameter param;
  setParameters(param);
  std::string model_path = "/home/daybeha/Documents/XYQ/yolov8_tensorrt/weights/yolov8n.trt";
  string dataset = "/media/daybeha/LTFM2/dataset/urban/urban39-pankyo/image/stereo_left";

  if(argc > 2){
    model_path = argv[1];
    dataset = argv[2];
  }
  // std::string video_path = "/home/daybeha/Documents/urban_split2.mp4";
  // std::string image_path = "/media/daybeha/LTFM2/dataset/kitti/sequences/00/image_0";
//    string dataset = "/media/daybeha/LTFM2/dataset/kitti/sequences/00";



//  setRenderWindow(param);
  YOLOV8 yolo(param);

  std::vector<unsigned char> trt_file = yolo_utils::loadModel(model_path);
  if (trt_file.empty())
  {
    sample::gLogError << "trt_file is empty!" << std::endl;
    return -1;
  }
  if (!yolo.init(trt_file))
  {
    sample::gLogError << "initEngine() ocur errors!" << std::endl;
    return -1;
  }
  yolo.check();

  // 读取图像路径
  vector<string> vstrImageLeft, vstrImageRight;
  vector<double> vTimestamps;
//    LoadKittiImages(dataset, vstrImageLeft, vstrImageRight, vTimestamps);
  LoadUrbanIamges(dataset, vstrImageLeft);
  const int nImages = vstrImageLeft.size();


  cv::Mat cur_img;
  std::vector<cv::Mat> imgs_batch;
  imgs_batch.reserve(param.batch_size);
  sample::gLogInfo << imgs_batch.capacity() << std::endl;
  for(int ni=0; ni<nImages; ni++)
  {
    cur_img = cv::imread(vstrImageLeft[ni], cv::IMREAD_COLOR); // 读取图像
    auto start = std::chrono::steady_clock::now();
    imgs_batch.emplace_back(cur_img);
    cv::imshow("original image", cur_img);

    task(yolo, param, imgs_batch, 1);

    imgs_batch.clear();
    auto end = std::chrono::steady_clock::now();
    auto time_usage = chrono::duration_cast<chrono::milliseconds>(end - start);
    std::cout << "total time: " << time_usage.count() << "ms\n";
    cv::waitKey(0);
  }

  return  -1;
}